package com.agent.vpsserver.modules.server.task;

import cn.hutool.core.date.DateTime;
import cn.hutool.core.date.DateUtil;
import com.agent.vpsserver.common.utils.RemoteCommandUtil;
import com.agent.vpsserver.modules.server.common.constants.CommandConstant;
import com.agent.vpsserver.modules.server.entity.pojo.ServerInfoDO;
import com.agent.vpsserver.modules.server.entity.pojo.UserAccountDO;
import com.agent.vpsserver.modules.server.entity.pojo.UserBandwidthDayDO;
import com.agent.vpsserver.modules.server.service.ServerInfoService;
import com.agent.vpsserver.modules.server.service.UserAccountService;
import com.agent.vpsserver.modules.server.service.UserBandwidthDayService;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.google.common.collect.Lists;
import com.jcraft.jsch.Session;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @Author: gc
 * @Date: 2024/3/13
 * @Description: 用户流量日统计
 **/
@Slf4j
@Component
@RequiredArgsConstructor
public class UserBandWidthDayTask {

    private final ServerInfoService serverInfoService;
    private final UserBandwidthDayService userBandwidthDayService;
    private final UserAccountService userAccountService;

    /**
     * 日统计数据-15分钟统计一次
     */
    @Scheduled(cron = "24 0/15 * * * ?")
    public void userBandWidthDayTask() {
        DateTime beginOfDay = DateUtil.beginOfDay(new Date());
        DateTime endOfDay = DateUtil.endOfDay(beginOfDay);
        String statisticsCommand = String.format(CommandConstant.STATISTICS_COMMAND, beginOfDay.getTime(), endOfDay.getTime());

        List<ServerInfoDO> serverInfos = serverInfoService.list();
        if (!CollectionUtils.isEmpty(serverInfos)) {
            for (ServerInfoDO serverInfo : serverInfos) {
                String remoteResponse = RemoteCommandUtil.exec(serverInfo.getHost(), serverInfo.getPort(), serverInfo.getUsername(), serverInfo.getPassword(), statisticsCommand);
                log.info("服务器ip：{}, 流量统计结果:{}", serverInfo.getHost(), remoteResponse);
                if (!StringUtils.isEmpty(remoteResponse)) {
                    String[] userBandWidths = remoteResponse.split("\\n");
                    for (String userBandWidth : userBandWidths) {
                        String[] userBandWidthSplit = userBandWidth.split(" ");
                        String username = userBandWidthSplit[0];
                        if (username.equals("-")) {
                            continue;
                        }
                        // 流量，单位byte
                        String bandWidthStr = userBandWidthSplit[1];
                        Long bandWidth = Long.valueOf(bandWidthStr);
                        UserAccountDO userAccount = userAccountService.getOne(new LambdaQueryWrapper<UserAccountDO>()
                                .eq(UserAccountDO::getUsername, username));
                        if (userAccount == null) {
                            continue;
                        }

                        // 查询日汇总数据
                        List<UserBandwidthDayDO> userBandwidthDays = userBandwidthDayService.list(new LambdaQueryWrapper<UserBandwidthDayDO>()
                                .eq(UserBandwidthDayDO::getUserId, userAccount.getId()));

                        // 修改日汇总数据
                        UserBandwidthDayDO userBandwidthDay = userBandwidthDays.stream()
                                .filter(info -> info.getDate().compareTo(beginOfDay) == 0)
                                .findFirst().orElse(null);
                        if (userBandwidthDay != null) {
                            userBandwidthDay.setBandwidth(bandWidth);
                            userBandwidthDayService.updateById(userBandwidthDay);
                        } else {
                            userBandwidthDay = new UserBandwidthDayDO();
                            userBandwidthDay.setUserId(userAccount.getId());
                            userBandwidthDay.setDate(beginOfDay);
                            userBandwidthDay.setBandwidth(bandWidth);
                            userBandwidthDayService.save(userBandwidthDay);
                            userBandwidthDays.add(userBandwidthDay);
                        }

                        // 没有设置最大流量则不计算用户流量
                        if (userAccount.getSurplusBandwidth() == null) {
                            continue;
                        }

                        // 计算用户流量是否超出，超出则封禁
                        long userBandWidthSum = userBandwidthDays.stream()
                                .mapToLong(UserBandwidthDayDO::getBandwidth)
                                .sum();
                        BigDecimal userBandWidthSumDecimal = new BigDecimal(String.valueOf(userBandWidthSum / 1024 / 1024 / 1024))
                                .setScale(2, RoundingMode.DOWN);
                        userAccount.setSurplusBandwidth(userAccount.getMaxBandwidth().subtract(userBandWidthSumDecimal));
                        if (userAccount.getSurplusBandwidth().compareTo(BigDecimal.ZERO) <= 0) {
                            userAccount.setStatus(2);
                            // 没有流量就移除账号
                            String removeUserCommand = String.format(CommandConstant.REMOVE_USER_COMMAND, username);
                            String removeUserRemoteResponse = RemoteCommandUtil.exec(serverInfo.getHost(), serverInfo.getPort(), serverInfo.getUsername(), serverInfo.getPassword(), removeUserCommand);
                            log.info("移除用户:{}，响应:{}", username, removeUserRemoteResponse);
                        }
                        userAccountService.updateById(userAccount);
                    }
                    String removeUserRemoteResponse = RemoteCommandUtil.exec(serverInfo.getHost(), serverInfo.getPort(), serverInfo.getUsername(), serverInfo.getPassword(), CommandConstant.SQUID_RECONFIGURE_COMMAND);
                    log.info("重新squid:{}，响应:{}", serverInfo.getHost(), removeUserRemoteResponse);
                }
            }
        }
    }

    /**
     * 日统计数据-统计昨天
     */
    @Scheduled(cron = "0 0 0 * * ?")
    public void userBandWidthYesterdayDayTask() {
        DateTime beginOfDay = DateUtil.beginOfDay(DateUtil.yesterday());
        DateTime endOfDay = DateUtil.endOfDay(beginOfDay);
        String statisticsCommand = String.format(CommandConstant.STATISTICS_COMMAND, beginOfDay.getTime() * 1000, endOfDay.getTime() * 1000);

        List<ServerInfoDO> serverInfos = serverInfoService.list();
        if (!CollectionUtils.isEmpty(serverInfos)) {
            for (ServerInfoDO serverInfo : serverInfos) {
                String remoteResponse = RemoteCommandUtil.exec(serverInfo.getHost(), serverInfo.getPort(), serverInfo.getUsername(), serverInfo.getPassword(), statisticsCommand);
                log.info("服务器ip：{}, 流量统计结果:{}", serverInfo.getHost(), remoteResponse);
                if (!StringUtils.isEmpty(remoteResponse)) {
                    String[] userBandWidths = remoteResponse.split("\\n");
                    for (String userBandWidth : userBandWidths) {
                        String[] userBandWidthSplit = userBandWidth.split(" ");
                        String username = userBandWidthSplit[0];
                        if (username.equals("-")) {
                            continue;
                        }
                        // 流量，单位byte
                        String bandWidthStr = userBandWidthSplit[1];
                        Long bandWidth = Long.valueOf(bandWidthStr);
                        UserAccountDO userAccount = userAccountService.getOne(new LambdaQueryWrapper<UserAccountDO>()
                                .eq(UserAccountDO::getUsername, username));
                        if (userAccount == null) {
                            continue;
                        }

                        // 修改日汇总数据
                        UserBandwidthDayDO userBandwidthDay = userBandwidthDayService.getOne(new LambdaQueryWrapper<UserBandwidthDayDO>()
                                .eq(UserBandwidthDayDO::getUserId, userAccount.getId()).
                                        eq(UserBandwidthDayDO::getDate, beginOfDay));
                        if (userBandwidthDay != null) {
                            userBandwidthDay.setBandwidth(bandWidth);
                            userBandwidthDayService.updateById(userBandwidthDay);
                        } else {
                            userBandwidthDay = new UserBandwidthDayDO();
                            userBandwidthDay.setUserId(userAccount.getId());
                            userBandwidthDay.setDate(beginOfDay);
                            userBandwidthDay.setBandwidth(bandWidth);
                            userBandwidthDayService.save(userBandwidthDay);
                        }
                    }
                    String removeUserRemoteResponse = RemoteCommandUtil.exec(serverInfo.getHost(), serverInfo.getPort(), serverInfo.getUsername(), serverInfo.getPassword(), CommandConstant.SQUID_RECONFIGURE_COMMAND);
                    log.info("重新squid:{}，响应:{}", serverInfo.getHost(), removeUserRemoteResponse);
                }
            }
        }
    }

    /**
     * 移除过期用户
     */
    @Scheduled(cron = "49 0/5 * * * ?")
    public void removeExpireUser() {
        // 过期用户
        List<UserAccountDO> accounts = userAccountService.list(new LambdaQueryWrapper<UserAccountDO>()
                .le(UserAccountDO::getExpireTime, new Date())
                .eq(UserAccountDO::getStatus, 1));
        if (CollectionUtils.isEmpty(accounts)) {
            return;
        }
        Map<Long, List<UserAccountDO>> serverIdMap = accounts.stream()
                .filter(userAccountDO -> userAccountDO.getExpireTime() != null)
                .collect(Collectors.groupingBy(UserAccountDO::getServerId));
        List<ServerInfoDO> serverInfos = serverInfoService.listByIds(serverIdMap.keySet());
        for (ServerInfoDO serverInfo : serverInfos) {
            Session session = RemoteCommandUtil.createSession(serverInfo.getHost(), serverInfo.getPort(), serverInfo.getUsername(), serverInfo.getPassword());
            try {
                List<UserAccountDO> userAccounts = serverIdMap.get(serverInfo.getId());
                List<List<UserAccountDO>> partitionUsers = Lists.partition(userAccounts, 50);
                for (List<UserAccountDO> partitionUser : partitionUsers) {
                    List<String> names = partitionUser.stream().map(UserAccountDO::getUsername).collect(Collectors.toList());
                    String removeUsername = String.join(" ", names);
                    String command = String.format(CommandConstant.REMOVE_USER_COMMAND, removeUsername);
                    String response = RemoteCommandUtil.exec(command, session);
                    log.info("执行命令:{}，结果:{}", command, response);
                }
            } finally {
                session.disconnect();
            }
        }
        userAccountService.update(new LambdaUpdateWrapper<UserAccountDO>()
                .in(UserAccountDO::getId, serverIdMap.keySet())
                .set(UserAccountDO::getStatus, 2));
    }
}
